#include "ObjectDetector.h"

#include "ImageUtility.h"

#include <QMetaType>
#include <QDebug>

using namespace cv;
#ifdef HAVE_CUDA
using namespace cv::gpu;
#endif

Q_DECLARE_METATYPE(QVector<QRect>)

ObjectDetector::ObjectDetector(QObject *parent) :
    QObject(parent),
    type(FACE_FRONTAL),
    shouldUseGPU(false),
    scaleFactor(1.0f),
    ROI(QRect()),
    windowStartScale(1.1),
    windowEndScale(10),
    windowSize(10, 10),
    cascadeCPU(NULL)
#ifdef HAVE_CUDA
    ,cascadeGPU(NULL)
#endif
{
    qRegisterMetaType<QVector<QRect> >();

    this->setupCascadeFilenames();
    this->setObjectType(type);
}

ObjectDetector::ObjectDetector(ObjectType type,
                               bool useGPU,
                               float imageScaleFactor,
                               QRect regionOfInterest,
                               float windowStartScaleFactor,
                               float windowEndScaleFactor,
                               int windowStartWidth,
                               int windowStartHeight) :
    QObject(),
    type(type),
    shouldUseGPU(useGPU),
    scaleFactor(imageScaleFactor),
    ROI(regionOfInterest),
    windowStartScale(windowStartScaleFactor),
    windowEndScale(windowEndScaleFactor),
    windowSize(windowStartWidth, windowStartHeight),
    cascadeCPU(NULL)
#ifdef HAVE_CUDA
    ,cascadeGPU(NULL)
#endif
{
    qRegisterMetaType<QVector<QRect> >();

    this->setupCascadeFilenames();
    this->setObjectType(type);

}

void ObjectDetector::setupCascadeFilenames()
{
    cascadeFilenames.clear();
    QString baseName = "/home/rivit/Demos/FaceRecognition/haarCascades/haarcascade_";

    cascadeFilenames.insert(FACE_FRONTAL, baseName + "frontalface_default.xml");
    cascadeFilenames.insert(FACE_FRONTAL_2, baseName + "frontalface_alt.xml");
    cascadeFilenames.insert(FACE_FRONTAL_3, baseName + "frontalface_alt2.xml");
    cascadeFilenames.insert(FACE_FRONTAL_4, baseName + "frontalface_alt_tree.xml");
    cascadeFilenames.insert(FACE_PROFILE, baseName + "profileface.xml");
    cascadeFilenames.insert(EYE, baseName + "eye.xml");
    cascadeFilenames.insert(EYE_WITH_GLASSES, baseName + "eye_tree_eyeglasses.xml");
    cascadeFilenames.insert(EYE_LEFT, baseName + "lefteye_2splits.xml");
    cascadeFilenames.insert(EYE_RIGHT, baseName + "righteye_2splits.xml");
    cascadeFilenames.insert(EYE_LEFT_ALT, baseName + "mcs_lefteye.xml");
    cascadeFilenames.insert(EYE_RIGHT_ALT, baseName + "mcs_righteye.xml");
    cascadeFilenames.insert(EYE_PAIR_BIG, baseName + "mcs_eyepair_big.xml");
    cascadeFilenames.insert(EYE_PAIR_SMALL, baseName + "mcs_eyepair_small.xml");
    cascadeFilenames.insert(BODY_UPPER, baseName + "upperbody.xml");
    cascadeFilenames.insert(BODY_UPPER_ALT, baseName + "mcs_upperbody.xml");
    cascadeFilenames.insert(BODY_LOWER, baseName + "lowerbody.xml");
    cascadeFilenames.insert(BODY_FULL, baseName + "fullbody.xml");
    cascadeFilenames.insert(MOUTH, baseName + "mcs_mouth.xml");
    cascadeFilenames.insert(NOSE, baseName + "mcs_nose.xml");
}

ObjectDetector::~ObjectDetector()
{

    if (cascadeCPU) {
        delete cascadeCPU;
    }

#ifdef HAVE_CUDA
    if (cascadeGPU) {
        delete cascadeGPU;
    }
#endif
}

ObjectDetector::ObjectType ObjectDetector::objectType()
{
    return type;
}

float ObjectDetector::imageScaleFactor()
{
    return scaleFactor;
}

QRect ObjectDetector::regionOfInterest()
{
    return ROI;
}

float ObjectDetector::windowStartScaleFactor()
{
    return windowStartScale;
}

float ObjectDetector::windowEndScaleFactor()
{
    return windowEndScale;
}

QSize ObjectDetector::windowStartSize()
{
    return windowSize;
}

void ObjectDetector::setObjectType(ObjectType objectType)
{
    type = objectType;
    if (cascadeCPU) {
        delete cascadeCPU;
    }

#ifdef HAVE_CUDA
    if (cascadeGPU) {
        delete cascadeGPU;
    }
#endif

    bool haveCuda = false;
#ifdef HAVE_CUDA
    haveCuda = true;
#endif

    if (shouldUseGPU && haveCuda) {
#ifdef HAVE_CUDA
        cascadeGPU = new cv::gpu::CascadeClassifier_GPU();
        cascadeGPU->load(cascadeFilenames.value(objectType).toStdString().c_str());
#endif
    } else {
        cascadeCPU = new cv::CascadeClassifier();
        cascadeCPU->load(cascadeFilenames.value(objectType).toStdString().c_str());
    }
}

void ObjectDetector::setImageScaleFactor(float scale)
{
    scaleFactor = scale;
}

void ObjectDetector::setRegionOfInterest(QRect rect)
{
    ROI = rect;
}

void ObjectDetector::setWindowStartScaleFactor(float start)
{
    windowStartScale = start;
}

void ObjectDetector::setWindowEndScaleFactor(float end)
{
    windowEndScale = end;
}

void ObjectDetector::setWindowStartSize(int width, int height)
{
    windowSize = QSize(width, height);
}

QVector<QVector<QRect> > ObjectDetector::detectObjectsInROIs(QVector<QRect> rects, QImage image)
{
    QVector<QVector<QRect> > objectList;
    foreach (QRect rect, rects) {
        this->setRegionOfInterest(rect);
        QVector<QRect> objects = this->detectObjects(image);
        objectList.append(objects);
    }

    return objectList;
}

QVector<QVector<QRect> > ObjectDetector::detectObjectsInROIs(QVector<QRect> rects, IplImage* image)
{
    QVector<QVector<QRect> > objectList;
    foreach (QRect rect, rects) {
        this->setRegionOfInterest(rect);
        QVector<QRect> objects = this->detectObjects(image);
        objectList.append(objects);
    }

    return objectList;
}

QVector<QRect> ObjectDetector::detectObjects(QImage image)
{
    IplImage* ipl = ImageUtility::QImageToIplImage(image);

    QVector<QRect> objects = this->detectObjects(Mat(ipl));

    emit detectedObjects(objects, image);

    cvReleaseImage(&ipl);

    return objects;
}

QVector<QRect> ObjectDetector::rectsToQRects(Rect* objects, int count)
{
    QVector<QRect> rects;
    if (count == 0)
        return rects;

    float reverseScale = 1.0f/scaleFactor;

    QPoint offset = ROI.topLeft();

    for (int i = 0; i < count; i++) {
        Rect rect = objects[i];

        QPoint topLeft(offset.x() + rect.x * reverseScale, offset.y() + rect.y * reverseScale);
        QSize size(rect.width * reverseScale, rect.height * reverseScale);

        rects.append(QRect(topLeft, size));
    }

    return rects;
}

QVector<QRect> ObjectDetector::detectObjectsGPU(Mat image)
{
#ifdef HAVE_CUDA
    GpuMat gpuImage, gray, scaled;

    gpuImage.upload(image);

    cv::gpu::cvtColor(gpuImage, gray, CV_BGR2GRAY);
    if (!ROI.isNull()) {
        cv::Size size = cascadeGPU->getClassifierSize();
        if (ROI.width() * scaleFactor > size.width && ROI.height() * scaleFactor > size.height)
            gray = Mat(gray, cv::Rect(ROI.x(), ROI.y(), ROI.width(), ROI.height()));
    }

    Size newSize(cvRound(gray.cols * scaleFactor), cvRound(gray.rows * scaleFactor));

    cv::gpu::resize(gray, scaled, newSize);

    Mat downloadedObjects;

    GpuMat gpuRects;

    int faceCount = cascadeGPU->detectMultiScale(scaled, gpuRects, windowStartScale, 4,
                                                  cv::Size(windowSize.width(), windowSize.height()));

    gpuRects.colRange(0, faceCount).download(downloadedObjects);

    Rect* objects = downloadedObjects.ptr<Rect>();

    return this->rectsToQRects(objects, faceCount);
#else
    return this->detectObjectsCPU(image);
#endif
}

QVector<QRect> ObjectDetector::detectObjectsCPU(Mat image)
{

    Mat gray, scaled;
    cvtColor(image, gray, CV_BGR2GRAY);
    if (!ROI.isNull()) {
        gray = Mat(gray, cv::Rect(ROI.x(), ROI.y(), ROI.width(), ROI.height()));
    }

    Size newSize(cvRound(gray.cols * scaleFactor), cvRound(gray.rows * scaleFactor));
    resize(gray, scaled, newSize);


    std::vector<Rect> cpuRects;

    cascadeCPU->detectMultiScale(scaled, cpuRects, windowStartScale, 4, CV_HAAR_DO_CANNY_PRUNING,
                                 cv::Size(windowSize.width(), windowSize.height()),
                                 cv::Size(windowSize.width() * windowEndScale, windowSize.height() * windowEndScale));


    if (cpuRects.empty()) {
        return QVector<QRect>();
    }
    Rect* objects = &cpuRects[0];

    return this->rectsToQRects(objects, cpuRects.size());
}

QVector<QRect> ObjectDetector::detectObjects(Mat image)
{
    if (shouldUseGPU) {
        return this->detectObjectsGPU(image);
    } else {
        return this->detectObjectsCPU(image);
    }
}
